import os
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
import numpy as np

class MyDataset(Dataset):
    def __init__(self, img_dir, label_dir, transform=None):
        """
        Args:
            img_dir (str): 图像文件（.jpg）所在的目录。
            label_dir (str): 标签文件（.png）所在的目录。
            transform (callable, optional): 一个可选的函数/变换，用于对样本进行预处理。
        """
        self.image_dir = img_dir
        self.label_dir = label_dir
        self.transform = transform

        # 获取所有图像的文件名（只考虑.jpg文件）
        self.image_filenames = [f for f in os.listdir(self.image_dir) if f.endswith('.jpg')]
        # 获取所有标签的文件名前缀（只考虑.png文件，并去掉扩展名）
        #label_filenames = [f.split('.')[0] for f in os.listdir(self.label_dir) if f.endswith('.png')]
        label_filenames = [f.split('.')[0] for f in os.listdir(self.label_dir) if f.endswith('.jpg')]
        # 确保图像和标签文件数量一致
        #assert len(self.image_filenames) == len(label_filenames), \
        #    "The number of images and labels must be the same."

    def __len__(self):
        # 数据集的大小是图像文件的数量
        return len(self.image_filenames)

    def __getitem__(self, idx):
        # 根据索引获取一个样本
        img_name = self.image_filenames[idx]
        img_path = os.path.join(self.image_dir, img_name)

        # 去掉图像文件名的扩展名，以获取对应的标签文件名
        #label_name = img_name.split('.')[0] + '.png'
        label_name = img_name.split('.')[0] + '.jpg'
        label_path = os.path.join(self.label_dir, label_name)

        # 加载图像和标签
        image = Image.open(img_path).convert("RGB")  # 图像是RGB三通道的
        label = Image.open(label_path).convert("L")  # 标签是单通道的（灰度图）

        # # 将PIL图像转换为NumPy数组
        # image_np = np.array(image)
        # label_np = np.array(label)

        # # 如果提供了变换函数，则对样本应用它
        # if self.transform:
        #     transformed = self.transform(image=image_np, mask=label_np)
        #     image_np = transformed['image']
        #     label_np = transformed['mask']
        #
        # # 将图像和标签转换为PyTorch张量
        # image_tensor = transforms.ToTensor()(image_np)
        # label_tensor = transforms.ToTensor()(label_np)
        #
        # # 返回图像和对应的标签
        # return image_tensor, label_tensor

        if self.transform:
                image = self.transform(image)
                label = self.transform(label)

        return image ,label